package midend.IR_Optimizer;

import com.sun.jdi.Value;
import midend.*;
import midend.IR_Type.IR_Int1;
import midend.IR_Type.IR_Int32;
import midend.IR_Type.IR_Type;
import midend.IR_Value.*;
import midend.IR_Instruction.*;
import midend.IR_Type.IR_Const;

import java.util.*;

public class LoopUnroll {
    private IR_Module module;
    private boolean change = false;
    private HashMap<IR_Value, IR_Value> valueMap = new HashMap<>();
    private IR_Builder builder = new IR_Builder();

    public LoopUnroll(IR_Module module) {
        this.module = module;
    }

    public void Optimize() {
        change = true;
        while (change) {
            change = false;
            for (IR_FuncBlock function : module.getFuncBlocks()) {
                if(function.get_bb_List().size()>3500){
                    continue;
                }
                LoopInfo loopInfo = new LoopInfo();
                Iterator_Analysis iteratorAnalysis = new Iterator_Analysis();
                loopInfo.loopAnalysis(function);
                iteratorAnalysis.iteratorAnalysis(function);
                ArrayList<Loop> loops = new ArrayList<>();
                for (Loop loop : function.getLoops()) {
                    if (!loops.contains(loop) && canBeUnroll(loop))  {
                        unrollLoop(loop, function);
                        loops.add(loop);
                    }
                }
            }
            new CFG_Builder(module).run();
        }
    }

    public boolean canBeUnroll(Loop loop) {
        if (!(loop.getChildLoops().isEmpty() && loop.getLoopBody().size() == 2)) {
            // 只对没有嵌套循环和基本块数量为2的循环进行展开
            return false;
        }
        // 如果循环结束值太大，循环展开代价太高，就不展开
        if (loop.getItEnd() instanceof IR_Int_Value && ((IR_Int_Value)loop.getItEnd()).get_Int_Value()>8000){
            return false;
        }
        valueMap = new HashMap<>();
        IR_BaseBlock head = loop.getHeader();
        IR_FuncBlock infunc = head.get_FuncBlock();

        for (IR_BaseBlock block : loop.getLoopBody()) {
            for (Instruction instruction : block.get_Instruction_List()) {
                // 不对有Alloc指令和Ret指令的循环进行展开
                if (!(instruction instanceof Branch_Instruction || instruction instanceof Store_Instruction || instruction instanceof Pointer_Instruction
                        || instruction instanceof Load_Instruction || instruction instanceof Calc_Instruction || instruction instanceof Cond_Instruction
                        || instruction instanceof Phi_Instruction || instruction instanceof Zext_Instruction || instruction instanceof Call_Instruction
                        || instruction instanceof Type_Trans_Instruction)) {
                    return false;
                }
            }
        }

        Branch_Instruction loopBranch = (Branch_Instruction) head.get_Instruction_List().getLast();
        if (loopBranch.get_op().size() == 1
                || loopBranch.get_op().get(0) instanceof IR_Int_Value && (((IR_Int_Value) loopBranch.get_op().get(0)).get_Int_Value() == 1
                || ((IR_Int_Value) loopBranch.get_op().get(0)).get_Int_Value() == 0)) {
            // 不对死循环展开
            return false;
        }

        for (Instruction instruction : head.get_Instruction_List()) {
            if (instruction instanceof Phi_Instruction) {
                if (((Phi_Instruction) instruction).get_ComingValues().size() != 2) {
                    return false;
                }
                for (Map.Entry<IR_BaseBlock, IR_Value> entry : ((Phi_Instruction) instruction).get_ComingValues().entrySet()) {
                    if (!loop.getLoopBody().contains(entry.getKey())) {
                        valueMap.put(instruction.get_op().get(0), entry.getValue());
                        // 第一次循环迭代时，phi指令结果的引用用setComing中迭代变量的初始值代替
                    }
                }
            }
        }

        HashSet<IR_Value> waitCheck = new HashSet<>();  // 等待确认是否为常量的值
        HashSet<IR_Value> needKnown = new HashSet<>();

        if (!(loopBranch.get_op().get(0) instanceof IR_Int_Value)) {
            // 需要知道iterEnd是否为常量
            needKnown.add(loopBranch.get_op().get(0));
        }

        while (needKnown.size() != 0) {
            IR_Value temp = needKnown.iterator().next();
            needKnown.remove(temp);
            waitCheck.add(temp);

            if (isParamOrGlobal(module, infunc, temp)){
                continue;
            }

            Instruction instruction = infunc.find_Instruction(temp);
            if (instruction == null)
            {
                continue;
            }
            for (IR_Value op : instruction.get_op()) {
                if (op == temp)
                {
                    continue;
                }
                if (op instanceof IR_Int_Value || op instanceof IR_Float_Value) {
                    continue;
                }
                if (!waitCheck.contains(op)) {
                    needKnown.add(op);
                }
            }
        }

        boolean same;
        while (!waitCheck.isEmpty()) {
            same = true;
            ArrayList<IR_Value> removeCheck = new ArrayList<>();
            for (IR_Value value : waitCheck) {
                boolean known = true;
                if (isParamOrGlobal(module, infunc, value)) {
                    known = false;
                } else {
                    Instruction instruction = module.find_Instruction_InModule(value);
                    if (instruction == null) {
                        continue;
                    }
                    for (IR_Value op : instruction.get_op()) {
                        if (op == value){
                            continue;
                        }
                        if (!((valueMap.containsKey(op) && (valueMap.get(op) instanceof IR_Int_Value || valueMap.get(op) instanceof IR_Float_Value))
                                || (op instanceof IR_Int_Value || op instanceof IR_Float_Value))) {
                            known = false;
                            break;
                        }
                    }
                }
                if (known) {
                    removeCheck.add(value);
                }
            }

            if (!removeCheck.isEmpty()) {
                same = false;
                for (IR_Value value : removeCheck) {
                    waitCheck.remove(value);
                }
            }
            if (same) {
                break;
            }
        }

        if (!waitCheck.isEmpty()) {
            return false;
        }
        return true;
    }


    public void unrollLoop(Loop loop, IR_FuncBlock function) {
        IR_BaseBlock unrollBlock = new IR_BaseBlock("unroll_block", function);
        ArrayList<Instruction> instructions = new ArrayList<>();
        for (IR_BaseBlock block : loop.getLoopBody()) {
            instructions.addAll(block.get_Instruction_List());
        }
        int hold = 0;
        boolean firstLoop = true;
        boolean unrollSuccess = false;
        while (hold < 10000) {
            HashMap<IR_Value, IR_Value> phiParr = new HashMap<>();
            for (Instruction instruction : instructions) {
                hold++;
                if (instruction instanceof Phi_Instruction) {
//                  loop:
//                  %x = phi i32 [ 0, %entry ], [ %x.next, %loop ]
//                  %x.next = add i32 %x, 1
//                  br i1 %cond, label %loop, label %exit
//                  第一次迭代（firstLoop = true）：
//                    跳过 Phi 指令的处理，直接处理 add 和 br 指令。
//                    valueMap 中记录 %x.next = 1（假设初始 %x = 0）。
//                  第二次迭代（firstLoop = false）：
//                    处理 Phi 指令时，发现 %x 的来源之一是 %loop（循环体），于是从 valueMap 获取 %x.next 的最新值（1）。
//                    将 %x = 1 记录到 phiParr，随后更新到 valueMap。
//                    继续处理 add 指令，计算 %x.next = 2。
                    if (!firstLoop){
                        for (Map.Entry<IR_BaseBlock, IR_Value> entry : ((Phi_Instruction) instruction).get_ComingValues().entrySet()) {
                            if (loop.getLoopBody().contains(entry.getKey())) {
                                // 记录phi指令结果的最新值
                                valueMap.put(instruction.get_op().get(0), valueMap.get(entry.getValue()));
                            }
                        }
                    }
                }
//                else {
//                    if (!phiParr.isEmpty()) {
//                        for (IR_Value value : phiParr.keySet()) {
//                            valueMap.put(value, phiParr.get(value));
//                        }
//                    }
//                    phiParr.clear();
//                }

                if (instruction instanceof Calc_Instruction) {
                    ArrayList<IR_Value> ops = new ArrayList<>();
                    IR_Value op = new IR_Local_Value(((IR_Local_Value)instruction.get_op().get(0)).get_Local_Identity(), instruction.get_op().get(0).get_Type());
                    ops.add(op);
                    for (int i = 1; i < instruction.get_op().size(); i++) {
                        op = instruction.get_op().get(i);
                        if (valueMap.containsKey(op)) {
                            op = (valueMap.get(op));
                        }
                        ops.add(op);
                    }
                    Calc_Instruction sameIns = new Calc_Instruction(ops, unrollBlock, ((Calc_Instruction) instruction).get_op_Type());
//                    unrollBlock.add_InstructionList(sameIns);
//                    valueMap.put(instruction.get_op().get(0), sameIns.get_op().get(0));
                    //TODO:sameBir.cmpType = instr.cmpType;
                    if (ops.get(0).get_Type() instanceof IR_Int32 && ops.get(1) instanceof IR_Int_Value && ops.get(2) instanceof IR_Int_Value) {
                        IR_Value constant_int = fold_Int32_Constant(((Calc_Instruction)sameIns).get_op_Type(), (IR_Int_Value)sameIns.get_op().get(1), (IR_Int_Value)sameIns.get_op().get(2));
                        valueMap.put(instruction.get_op().get(0), constant_int);
                        sameIns.remove();
                    } else {
                        unrollBlock.add_InstructionList(sameIns);
                        valueMap.put(instruction.get_op().get(0), sameIns.get_op().get(0));
                    }
                }

                if (instruction instanceof Cond_Instruction){
                    ArrayList<IR_Value> ops = new ArrayList<>();
                    IR_Value op = new IR_Local_Value(((IR_Local_Value)instruction.get_op().get(0)).get_Local_Identity(), instruction.get_op().get(0).get_Type());
                    ops.add(op);
                    for (int i = 1; i < instruction.get_op().size(); i++) {
                        op = instruction.get_op().get(i);
                        if (valueMap.containsKey(op)) {
                            op = (valueMap.get(op));
                        }
                        ops.add(op);
                    }
                    Cond_Instruction sameIns = new Cond_Instruction(ops, unrollBlock, ((Cond_Instruction)instruction).get_cond());
                    if (ops.get(0).get_Type() instanceof IR_Int1 && ops.get(1) instanceof IR_Int_Value && ops.get(2) instanceof IR_Int_Value) {
                        IR_Value constant_int = fold_Int1_Constant(sameIns.get_cond(), (IR_Int_Value)sameIns.get_op().get(1), (IR_Int_Value)sameIns.get_op().get(2));
                        valueMap.put(instruction.get_op().get(0), constant_int);
                        sameIns.remove();
                    } else {
                        unrollBlock.add_InstructionList(sameIns);
                        valueMap.put(instruction.get_op().get(0), sameIns.get_op().get(0));
                    }
                }

                if (instruction instanceof Pointer_Instruction) {
                    ArrayList<IR_Value> ops = new ArrayList<>();
                    IR_Value op = new IR_Local_Value(((IR_Local_Value)instruction.get_op().get(0)).get_Local_Identity(), instruction.get_op().get(0).get_Type());
                    ops.add(op);
                    for (int i = 1; i < instruction.get_op().size(); i++) {
                        op = instruction.get_op().get(i);
                        if (valueMap.containsKey(op)) {
                            op = (valueMap.get(op));
                        }
                        ops.add(op);
                    }
                    Pointer_Instruction sameIns = new Pointer_Instruction(ops, unrollBlock);
                    unrollBlock.add_InstructionList(sameIns);
                    valueMap.put(instruction.get_op().get(0), sameIns.get_op().get(0));
                }

                if (instruction instanceof Load_Instruction) {
                    ArrayList<IR_Value> ops = new ArrayList<>();
                    IR_Value op = new IR_Local_Value(((IR_Local_Value)instruction.get_op().get(0)).get_Local_Identity(), instruction.get_op().get(0).get_Type());
                    ops.add(op);
                    for (int i = 1; i < instruction.get_op().size(); i++) {
                        op = instruction.get_op().get(i);
                        if (valueMap.containsKey(op)) {
                            op = (valueMap.get(op));
                        }
                        ops.add(op);
                    }
                    Load_Instruction sameIns = new Load_Instruction(ops, unrollBlock);
                    unrollBlock.add_InstructionList(sameIns);
                    valueMap.put(instruction.get_op().get(0), sameIns.get_op().get(0));
                }

                if (instruction instanceof Store_Instruction) {
                    ArrayList<IR_Value> ops = new ArrayList<>();
                    IR_Value op = new IR_Local_Value(((IR_Local_Value)instruction.get_op().get(0)).get_Local_Identity(), instruction.get_op().get(0).get_Type());
                    ops.add(op);
                    for (int i = 1; i < instruction.get_op().size(); i++) {
                        op = instruction.get_op().get(i);
                        if (valueMap.containsKey(op)) {
                            op = (valueMap.get(op));
                        }
                        ops.add(op);
                    }
                    Store_Instruction sameIns = new Store_Instruction(ops, unrollBlock);
                    unrollBlock.add_InstructionList(sameIns);
                    valueMap.put(instruction.get_op().get(0), sameIns.get_op().get(0));
                }

                if (instruction instanceof Zext_Instruction) {
                    ArrayList<IR_Value> ops = new ArrayList<>();
                    IR_Value op = new IR_Local_Value(((IR_Local_Value)instruction.get_op().get(0)).get_Local_Identity(), instruction.get_op().get(0).get_Type());
                    ops.add(op);
                    for (int i = 1; i < instruction.get_op().size(); i++) {
                        op = instruction.get_op().get(i);
                        if (valueMap.containsKey(op)) {
                            op = (valueMap.get(op));
                        }
                        ops.add(op);
                    }
                    Zext_Instruction sameIns = new Zext_Instruction(ops, unrollBlock);
                    unrollBlock.add_InstructionList(sameIns);
                    valueMap.put(instruction.get_op().get(0), sameIns.get_op().get(0));
                }

                if (instruction instanceof Call_Instruction) {
                    ArrayList<IR_Value> ops = new ArrayList<>();
                    IR_Value op = new IR_Local_Value(((IR_Local_Value)instruction.get_op().get(0)).get_Local_Identity(), instruction.get_op().get(0).get_Type());
                    ops.add(op);
                    for (int i = 1; i < instruction.get_op().size(); i++) {
                        op = instruction.get_op().get(i);
                        if (valueMap.containsKey(op)) {
                            op = (valueMap.get(op));
                        }
                        ops.add(op);
                    }
                    Call_Instruction sameIns = new Call_Instruction(ops, unrollBlock);
                    unrollBlock.add_InstructionList(sameIns);
                    valueMap.put(instruction.get_op().get(0), sameIns.get_op().get(0));
                }

                if (instruction instanceof Type_Trans_Instruction){
                    ArrayList<IR_Value> ops = new ArrayList<>();
                    IR_Value op = new IR_Local_Value(((IR_Local_Value)instruction.get_op().get(0)).get_Local_Identity(), instruction.get_op().get(0).get_Type());
                    ops.add(op);
                    for (int i = 1; i < instruction.get_op().size(); i++) {
                        op = instruction.get_op().get(i);
                        if (valueMap.containsKey(op)) {
                            op = (valueMap.get(op));
                        }
                        ops.add(op);
                    }
                    Type_Trans_Instruction sameIns = new Type_Trans_Instruction(ops, unrollBlock, ((Type_Trans_Instruction)instruction).get_Trans_Type());
                    unrollBlock.add_InstructionList(sameIns);
                    valueMap.put(instruction.get_op().get(0), sameIns.get_op().get(0));
                }

                if (instruction instanceof Branch_Instruction) {
                    if (((Branch_Instruction) instruction).br_type()) {
                        IR_Value cond = ((Branch_Instruction) instruction).get_op().get(0);
                        if (((IR_Int_Value) valueMap.get(cond)).get_Int_Value() == 0) {
//                            如果条件计算为 0（即 false），表示循环终止。
//                            此时设置 unrollSuccess = true，表示展开成功，可以替换原循环
                            unrollSuccess = true;
                            //替换block
                            for (IR_BaseBlock basicBlock : function.get_bb_List()) {
                                if (basicBlock.get_Instruction_List().getLast() instanceof Branch_Instruction) {
                                    ((Branch_Instruction) basicBlock.get_Instruction_List().getLast()).change_Dest(loop.getHeader(), unrollBlock);
                                }
                                if (basicBlock.get_Instruction_List().getFirst() instanceof Phi_Instruction) {
                                    Instruction instr = basicBlock.get_Instruction_List().getFirst();
                                    int count = 0;
                                    while (instr instanceof Phi_Instruction) {
                                        ((Phi_Instruction) instr).change_Origin(loop.getHeader(), unrollBlock);
                                        instr = basicBlock.get_Instruction_List().get(++count);
                                    }
                                }
                            }

                            for (IR_Value oldValue : valueMap.keySet()) {
                                function.replaceUseInFunc(oldValue, valueMap.get(oldValue));
                                Instruction oldIns = function.find_Instruction(oldValue);
                                if (oldIns != null) {
                                    oldIns.remove();
                                }
                            }
                            // 在unrollBlock最后加上跳转到循环出口的语句
                            Branch_Instruction brInstr = new Branch_Instruction(Collections.singletonList(instruction.get_op().get(2)), unrollBlock);
                            unrollBlock.add_InstructionList(brInstr);
                            int count = function.get_bb_List().indexOf(loop.getHeader());
                            function.get_bb_List().add(count, unrollBlock);
                            function.get_bb_List().removeAll(loop.getLoopBody());
                            break;
                        }
                    }
                }
            }

            firstLoop = false;
            if (unrollSuccess) {
                break;
            }
        }

        if (!unrollSuccess) {
            // 如果循环展开失败，进行回滚操作
            for (int count = unrollBlock.get_Instruction_List().size() - 1; count >= 0; count--) {
                Instruction instruction = unrollBlock.get_Instruction_List().get(count);
                instruction.remove();
            }
            return;
        }
        for (Instruction instruction : instructions) {
            instruction.remove();
        }
        change = true;
    }

    public boolean isParamOrGlobal(IR_Module module, IR_FuncBlock funcBlock, IR_Value value){
        if(module.get_Gloabal_Varmap().containsKey(value.get_ValueName()))
        {
            return true;
        }
        boolean isparaem = false;
        for (IR_Value param : funcBlock.get_params()){
            if (value == param) {
                return true;
            }
        }
        return false;
    }

    //计算32位整数常量表达式结果
    private IR_Value fold_Int32_Constant(String type, IR_Value lop, IR_Value rop) {
        IR_Type res_Type = lop.get_Type();
        if (res_Type instanceof IR_Int32) {
            int lop_Value = Integer.parseInt(lop.get_ValueName());
            int rop_Value = Integer.parseInt(rop.get_ValueName());
            int res_Value = 0;
            switch (type) {
                case "add":
                    res_Value = lop_Value + rop_Value;
                    break;
                case "sub":
                    res_Value = lop_Value - rop_Value;
                    break;
                case "mul":
                    res_Value = lop_Value * rop_Value;
                    break;
                case "sdiv":
                    res_Value = lop_Value / rop_Value;
                    break;
                case "srem":
                    res_Value = lop_Value % rop_Value;
                    break;
                case "xor":
                    res_Value = lop_Value ^ rop_Value;
                    break;
                case "shl":
                    res_Value = lop_Value << rop_Value;
            }
            return new IR_Int_Value(res_Value, IR_Int32.IR_Int32());
        }
        return null;
    }

    //计算常量表达式结果
    private IR_Value fold_Int1_Constant(int icmp, IR_Value lop, IR_Value rop) {
        IR_Type res_Type = lop.get_Type();
        if (res_Type instanceof IR_Int32) {
            int lop_Value = Integer.parseInt(lop.get_ValueName());
            int rop_Value = Integer.parseInt(rop.get_ValueName());
            boolean res_Value = false;
            switch (IR_Const.int_Compare[icmp]) {
                case "eq":
                    res_Value = lop_Value == rop_Value;
                    break;
                case "ne":
                    res_Value = lop_Value != rop_Value;
                    break;
                case "sgt":
                    res_Value = lop_Value > rop_Value;
                    break;
                case "sge":
                    res_Value = lop_Value >= rop_Value;
                    break;
                case "slt":
                    res_Value = lop_Value < rop_Value;
                    break;
                case "sle":
                    res_Value = lop_Value <= rop_Value;
                    break;
            }
            int res_int = res_Value ? 1 : 0;
            return new IR_Int_Value(res_int, IR_Int1.IR_Int1());
        }
        return null;
    }

}
